"""

【注意：这是实验代码。】

"""
from langchain.llms import OpenAI
from langchain.prompts import PromptTemplate
import os
import re
from PyCmpltrtok.util_mongo import get_history
from common import RESULT_PREFIX

NO_SYMPTOM = '没有症状'
COMMAND = f"根据医生和患者的对话，用逗号分隔的形式列举出患者的症状。\
不允许回答用户的提问，不允许在答案中添加编造成分，不允许在答案中生成新的闲聊，答案请使用中文。 如果没有闲聊内容或者在闲聊中患者没有陈述症状，请生成 “{NO_SYMPTOM}”。"
OK_AND_MORE = '了解了，还有别的症状吗？'
TEMPLATE = f"""<指令>{COMMAND}</指令>

例子:
<对话>
患者: 我老是症状A
医生: {OK_AND_MORE}
患者: 我还症状B
医生: OK
</对话>
<症状>症状A,症状B</症状>

当前对话:
<对话>
{{dialog}}
</对话>
<症状>"""


def analyze_symptom(xinput, xlog, model):
    # print(xlog)
    
    xstr = ''
    for xpair in xlog:
        xstr += '患者: ' + xpair[0].replace('\n\n\n', '\n').replace('\n\n', '\n') + '\n'
        xdoctor_reply = xpair[1].replace('\n\n\n', '\n').replace('\n\n', '\n')
        if xdoctor_reply[:len(RESULT_PREFIX)] == RESULT_PREFIX:
            xdoctor_reply = OK_AND_MORE
        xstr += '医生: ' + xdoctor_reply + '\n'
    xstr += '患者: ' + xinput + '\n'
    xstr += '医生: OK'
    
    prompt = PromptTemplate.from_template(TEMPLATE)
    chain = prompt | model
    xres = chain.invoke({
        'dialog': xstr,
    })
    print('RESULT:', xres)
    
    xres = xres.replace('<症状>', '')
    xres = xres.replace('</症状>', '')
    xres = xres.replace('<例子>', '')
    xres = xres.replace('</例子>', '')
    xres = xres.replace('<列表>', '')
    xres = xres.replace('</列表>', '')
    xres = xres.replace('\n', '')
    xres = xres.strip()
    
    symptoms = []
    xlist = re.split(r'[,，、]\s*', xres)
    for xelement in xlist:
        xelement = xelement.strip()
        if not xelement:
            continue
        if NO_SYMPTOM not in xelement:
            symptoms.append(xelement)
    return symptoms


def analyze_symptom4user(mdb, xinput, username, model):
    xlog = get_history(mdb, username)
    xres = analyze_symptom(xinput, xlog, model)
    return xres


if '__main__' == __name__:
    import pymongo as pm
    from PyCmpltrtok.common import sep
    from python_nlp.wechat.x00350_chatglm2_redis_history.daemon import MONGODB_NAME
    
    username = 'oW1QH6VuJHlpGugb0XrzlNEV2wBQ'
    
    # 连接Mongodb
    mongo = pm.MongoClient(os.environ['M_HOST'], int(os.environ['M_PORT']), serverSelectionTimeoutMS=3000)
    mdb = mongo[MONGODB_NAME]
    get_history(mdb, 'user_xxxx')  # try it

    model = OpenAI(
        streaming=True,
        verbose=True,
        callbacks=[],
        openai_api_key="token1",
        openai_api_base=f"http://127.0.0.1:3000/v1",
        model_name="chatglm2-6b",
        temperature=0.0,
        openai_proxy=None,
        top_p=1.0,
    )

    xinput = '我还经常打喷嚏'
    xres = analyze_symptom4user(mdb, xinput, username, model)
    sep('RESULT')
    print(xres)
    